package com.hexiang.servlet;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.DispatcherServlet;
import org.springframework.web.util.ContentCachingRequestWrapper;
import org.springframework.web.util.ContentCachingResponseWrapper;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Collection;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

@Slf4j
public class LoggableDispatcherServlet extends DispatcherServlet {


    private static final ObjectMapper mapper = new ObjectMapper();

    @Override
    protected void doDispatch(HttpServletRequest request, HttpServletResponse response) throws Exception {
        ContentCachingRequestWrapper requestWrapper = new ContentCachingRequestWrapper(request);
        ContentCachingResponseWrapper responseWrapper = new ContentCachingResponseWrapper(response);
        //创建一个 json 对象，用来存放 http 日志信息
        Map<String, Object> logMap = new HashMap<>();
        logMap.put("uri", requestWrapper.getRequestURI());
        logMap.put("clientIp", requestWrapper.getRemoteAddr());
        logMap.put("requestHeaders", getRequestHeaders(requestWrapper));
        String method = requestWrapper.getMethod();
        logMap.put("method", method);
        try {
            super.doDispatch(requestWrapper, responseWrapper);
        } finally {
            if (method.equals("GET")) {
                logMap.put("request", requestWrapper.getParameterMap());
            } else {
                logMap.put("request", new String(requestWrapper.getContentAsByteArray()));
            }
            logMap.put("status", responseWrapper.getStatus());
            logMap.put("response", new String(responseWrapper.getContentAsByteArray()));
            responseWrapper.copyBodyToResponse();

            logMap.put("responseHeaders", getResponseHeaders(responseWrapper));
            log.info(logMap.toString());
        }
    }

    private Map<String, Object> getRequestHeaders(HttpServletRequest request) {
        Map<String, Object> headers = new HashMap<>();
        Enumeration<String> headerNames = request.getHeaderNames();
        while (headerNames.hasMoreElements()) {
            String headerName = headerNames.nextElement();
            headers.put(headerName, request.getHeader(headerName));
        }
        return headers;

    }

    private Map<String, Object> getResponseHeaders(ContentCachingResponseWrapper response) {
        Map<String, Object> headers = new HashMap<>();
        Collection<String> headerNames = response.getHeaderNames();
        for (String headerName : headerNames) {
            headers.put(headerName, response.getHeader(headerName));
        }
        return headers;
    }
}